# -*- coding: utf-8 -*-
"""
MindMate AI聊天服务器 - 带豆包AI集成
"""

from fastapi import FastAPI, HTTPException, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
from datetime import datetime
import sys
import os
import json
import asyncio
import shutil
from pathlib import Path

# 导入模型管理器
from model_manager import model_manager

# 导入搜索服务
from tavily_search_service import search_service

# 导入数据库服务
from database import get_db_service

app = FastAPI(
    title="MindMate AI Chat Backend",
    description="AI Learning Companion System with Doubao Integration",
    version="1.0.0"
)

# 添加CORS支持
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000", "http://127.0.0.1:3000", "http://localhost:3002", "http://127.0.0.1:3002"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 添加验证错误处理器
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    print(f"验证错误 - URL: {request.url}")
    print(f"错误详情: {exc.errors()}")
    try:
        body = await request.body()
        print(f"请求体: {body.decode()}")
    except:
        print("无法读取请求体")
    return JSONResponse(
        status_code=422,
        content={"detail": exc.errors()}
    )

# 使用全局模型管理器（已在model_manager.py中初始化）

# 请求数据模型
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[ChatMessage]
    model: str = "doubao-seed-1-6-250615"  # 默认模型
    searchEnabled: bool = False  # 是否启用搜索增强
    sessionId: Optional[str] = None  # 会话ID，用于获取上下文

class ChatResponse(BaseModel):
    content: str
    model: str
    timestamp: str

@app.get("/")
async def root():
    return {
        "message": "MindMate AI Chat Backend is running!",
        "status": "healthy",
        "python_version": sys.version.split()[0],
        "timestamp": datetime.now().isoformat(),
        "features": [
            "豆包AI集成",
            "聊天API", 
            "健康检查"
        ],
        "endpoints": {
            "chat": "POST /api/chat",
            "models": "GET /api/models",
            "test_ai": "GET /test/ai",
            "docs": "GET /docs (API文档)",
            "openapi": "GET /openapi.json"
        }
    }

@app.get("/test-ui")
async def test_ui():
    """返回测试界面"""
    return FileResponse("test_ui.html")

@app.get("/health")
async def health():
    return {"status": "ok", "timestamp": datetime.now().isoformat()}

@app.post("/api/chat", response_model=ChatResponse)
async def chat_completion(request: ChatRequest):
    """AI聊天接口"""
    try:
        print(f"收到聊天请求: model={request.model}, messages数量={len(request.messages)}, searchEnabled={request.searchEnabled}, sessionId={request.sessionId}")
        for i, msg in enumerate(request.messages):
            print(f"消息{i}: role={msg.role}, content长度={len(msg.content)}")
        
        # 如果提供了sessionId，使用会话上下文（包括备忘录）
        if request.sessionId:
            # 优先使用数据库服务
            if db_service and db_service.initialize():
                session_data = await db_service.get_session(request.sessionId)
                if session_data:
                    # 转换为SessionData格式
                    session = SessionData(
                        id=session_data.id,
                        title=session_data.title,
                        roleId=session_data.role_id,
                        model=session_data.model,
                        searchEnabled=session_data.search_enabled,
                        memo=session_data.memo,
                        memoEnabled=session_data.memo_enabled,
                        isFavorite=session_data.is_favorite,
                        isArchived=session_data.is_archived,
                        createdAt=session_data.created_at,
                        updatedAt=session_data.updated_at,
                        messages=[]
                    )
                else:
                    session = None
            else:
                # 降级到内存存储
                session = await get_session_by_id(request.sessionId)
            if session:
                # 使用会话上下文消息（包括备忘录）
                messages = get_context_messages(session)
                print(f"使用会话上下文: 会话ID={request.sessionId}, 上下文消息数量={len(messages)}")
                if session.memo:
                    print(f"包含备忘录: {session.memo[:100]}...")
            else:
                print(f"会话未找到: {request.sessionId}, 使用原始消息")
                messages = [
                    {"role": msg.role, "content": msg.content}
                    for msg in request.messages
                ]
        else:
            # 转换消息格式
            messages = [
                {"role": msg.role, "content": msg.content}
                for msg in request.messages
            ]
        
        # 如果启用搜索功能，处理最新的用户消息
        if request.searchEnabled and messages and messages[-1]["role"] == "user":
            user_query = messages[-1]["content"]
            
            # 判断是否需要搜索
            if search_service.should_trigger_search(user_query):
                print(f"触发搜索: 查询='{user_query}'")
                
                # 执行搜索
                search_results = await search_service.search(user_query)
                
                if search_results:
                    # 格式化搜索结果为上下文
                    search_context = search_service.format_results_for_context(search_results)
                    
                    # 在用户消息前添加搜索上下文
                    enhanced_user_message = f"{search_context}\n\n## 用户问题\n{user_query}\n\n请基于以上搜索结果和你的知识来回答用户的问题。如果搜索结果相关，请优先参考搜索结果中的信息。"
                    
                    # 更新最后一条用户消息
                    messages[-1]["content"] = enhanced_user_message
                    
                    print(f"搜索增强完成: 找到{len(search_results)}个结果")
                else:
                    print("搜索未找到相关结果，使用原始查询")
            else:
                print("查询不需要搜索增强")
        
        # 使用ModelManager调用指定模型
        response_content = await model_manager.chat_completion(messages, request.model)
        
        return ChatResponse(
            content=response_content,
            model=request.model,
            timestamp=datetime.now().isoformat()
        )
        
    except Exception as e:
        print(f"聊天API异常: {str(e)}")
        raise HTTPException(status_code=500, detail=f"AI service error: {str(e)}")

@app.get("/api/models")
async def get_models():
    """获取可用模型列表"""
    try:
        models = model_manager.get_available_models()
        return {
            "models": models,
            "current": model_manager.get_current_model(),
            "default": "doubao-seed-1-6-250615"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取模型列表失败: {str(e)}")

@app.get("/api/search/test")
async def test_search():
    """测试搜索功能"""
    try:
        # 测试查询
        test_query = "2024年最新的AI技术发展"
        
        print(f"测试搜索功能: 查询='{test_query}'")
        
        # 检查是否应该触发搜索
        should_search = search_service.should_trigger_search(test_query)
        print(f"是否应该触发搜索: {should_search}")
        
        if should_search:
            # 执行搜索
            results = await search_service.search(test_query)
            print(f"搜索结果数量: {len(results)}")
            
            # 格式化结果
            formatted_context = search_service.format_results_for_context(results)
            
            return {
                "success": True,
                "query": test_query,
                "should_trigger_search": should_search,
                "results_count": len(results),
                "results": [{"title": r.title, "url": r.url, "snippet": r.snippet[:100]} for r in results],
                "formatted_context": formatted_context[:500] + "..." if len(formatted_context) > 500 else formatted_context
            }
        else:
            return {
                "success": True,
                "query": test_query,
                "should_trigger_search": should_search,
                "message": "查询不会触发搜索"
            }
            
    except Exception as e:
        print(f"搜索测试异常: {str(e)}")
        return {
            "success": False,
            "error": str(e)
        }

@app.get("/api/api-status")
async def get_api_status():
    """获取API状态"""
    try:
        # 测试当前模型连接
        current_model = model_manager.get_current_model()
        test_result = model_manager.test_model_connection(current_model)
        
        return {
            "status": "healthy" if test_result["success"] else "unhealthy",
            "ai_service": "connected" if test_result["success"] else "disconnected",
            "current_model": current_model,
            "timestamp": datetime.now().isoformat(),
            "server": "MindMate AI Chat Server",
            "version": "1.0.0"
        }
    except Exception as e:
        return {
            "status": "error",
            "ai_service": "error",
            "error": str(e),
            "timestamp": datetime.now().isoformat()
        }

@app.get("/test/ai")
async def test_ai():
    """测试AI功能"""
    try:
        current_model = model_manager.get_current_model()
        test_result = model_manager.test_model_connection(current_model)
        
        if test_result["success"]:
            test_messages = [
                {"role": "user", "content": "你好！"}
            ]
            
            response = await model_manager.chat_completion(test_messages, current_model)
            
            return {
                "status": "success",
                "message": f"AI service is working (Model: {current_model})",
                "model": current_model,
                "test_response": response[:100] + "..." if len(response) > 100 else response
            }
        else:
            return {
                "status": "failed", 
                "message": f"AI service connection failed: {test_result['message']}",
                "model": current_model
            }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

@app.post("/api/models/switch")
async def switch_model(request: dict):
    """切换当前使用的模型"""
    try:
        model_id = request.get("model_id")
        if not model_id:
            raise HTTPException(status_code=400, detail="模型ID不能为空")
        
        success = model_manager.set_current_model(model_id)
        if success:
            return {
                "success": True,
                "message": f"已切换到模型: {model_id}",
                "current_model": model_manager.get_current_model()
            }
        else:
            return {
                "success": False,
                "message": f"切换模型失败: 模型 {model_id} 不可用或未配置"
            }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"切换模型失败: {str(e)}")

@app.post("/api/models/configure")
async def configure_model(request: dict):
    """配置模型API密钥"""
    try:
        model_id = request.get("model_id")
        api_key = request.get("api_key")
        
        if not model_id or not api_key:
            raise HTTPException(status_code=400, detail="模型ID和API密钥不能为空")
        
        result = model_manager.configure_model(model_id, api_key)
        
        if result["success"]:
            return {
                "success": True,
                "message": result["message"],
                "model_id": model_id
            }
        else:
            return {
                "success": False,
                "message": result["message"],
                "model_id": model_id
            }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"配置模型失败: {str(e)}")

@app.get("/api/models/test/{model_id}")
async def test_model(model_id: str):
    """测试指定模型的连接"""
    try:
        result = model_manager.test_model_connection(model_id)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"测试模型连接失败: {str(e)}")

# 数据文件路径配置
DATA_DIR = Path("data")
SESSIONS_FILE = DATA_DIR / "sessions.json"
ROLES_FILE = DATA_DIR / "roles.json"

# 确保数据目录存在
DATA_DIR.mkdir(exist_ok=True)

# 数据库服务初始化
db_service = None

# 兼容性：保留内存存储作为备份
sessions_db = {}  # 内存存储
session_counter = 0

# 自定义角色管理
roles_db = {}  # 内存存储角色
role_counter = 0

# 数据持久化功能
def save_sessions():
    """保存会话数据到文件"""
    try:
        # 先保存到临时文件，然后重命名，确保原子性操作
        temp_file = SESSIONS_FILE.with_suffix('.tmp')
        
        # 将 Pydantic 模型正确序列化为字典
        serialized_sessions = {}
        for session_id, session in sessions_db.items():
            serialized_sessions[session_id] = {
                "id": session.id,
                "title": session.title,
                "lastActiveAt": session.lastActiveAt,
                "lastMessagePreview": session.lastMessagePreview,
                "isFavorite": session.isFavorite,
                "isArchived": session.isArchived,
                "roleId": session.roleId,
                "model": session.model,
                "searchEnabled": session.searchEnabled,
                "messages": session.messages,  # 消息列表已经是字典格式
                "memo": session.memo,
                "memo_created_at": session.memo_created_at
            }
        
        with open(temp_file, 'w', encoding='utf-8') as f:
            json.dump(serialized_sessions, f, ensure_ascii=False, indent=2)
        temp_file.replace(SESSIONS_FILE)
        print(f"会话数据已保存到 {SESSIONS_FILE}")
    except Exception as e:
        print(f"保存会话数据失败: {e}")
        # 尝试从备份恢复
        if SESSIONS_FILE.exists():
            print("尝试从现有文件继续运行")
        else:
            print("创建空的会话数据文件")
            SESSIONS_FILE.write_text('{}', encoding='utf-8')

def save_roles():
    """保存角色数据到文件"""
    try:
        # 先保存到临时文件，然后重命名，确保原子性操作
        temp_file = ROLES_FILE.with_suffix('.tmp')
        
        # 将 Pydantic 模型正确序列化为字典
        serialized_roles = {}
        for role_id, role in roles_db.items():
            serialized_roles[role_id] = {
                "id": role.id,
                "name": role.name,
                "title": role.title or "",
                "input_spec": role.input_spec,
                "output_spec": role.output_spec,
                "custom_traits": [{"key": trait.key, "value": trait.value} for trait in role.custom_traits] if role.custom_traits else [],
                "system_prompt_snapshot": role.system_prompt_snapshot,
                "is_builtin": role.is_builtin,
                "created_at": role.created_at,
                "updated_at": role.updated_at,
                "usage_count": role.usage_count,
                "last_used_at": role.last_used_at
            }
        
        with open(temp_file, 'w', encoding='utf-8') as f:
            json.dump(serialized_roles, f, ensure_ascii=False, indent=2)
        temp_file.replace(ROLES_FILE)
        print(f"角色数据已保存到 {ROLES_FILE}")
    except Exception as e:
        print(f"保存角色数据失败: {e}")
        # 尝试从备份恢复
        if ROLES_FILE.exists():
            print("尝试从现有文件继续运行")
        else:
            print("创建空的角色数据文件")
            ROLES_FILE.write_text('{}', encoding='utf-8')

def load_sessions():
    """从文件加载会话数据"""
    global sessions_db, session_counter
    try:
        if SESSIONS_FILE.exists():
            with open(SESSIONS_FILE, 'r', encoding='utf-8') as f:
                loaded_data = json.load(f)
                
                # 将 JSON 数据反序列化为 Pydantic 模型
                for session_id, session_data in loaded_data.items():
                    # 如果数据是字符串（旧格式），跳过
                    if isinstance(session_data, str):
                        print(f"跳过旧格式会话数据: {session_id}")
                        continue
                    
                    # 创建 SessionData 对象
                    session = SessionData(
                        id=session_data["id"],
                        title=session_data["title"],
                        lastActiveAt=session_data["lastActiveAt"],
                        lastMessagePreview=session_data.get("lastMessagePreview", ""),
                        isFavorite=session_data.get("isFavorite", False),
                        isArchived=session_data.get("isArchived", False),
                        roleId=session_data.get("roleId", "emma"),
                        model=session_data.get("model", "doubao-seed-1-6-250615"),
                        searchEnabled=session_data.get("searchEnabled", True),
                        messages=session_data.get("messages", []),
                        memo=session_data.get("memo", ""),
                        memo_created_at=session_data.get("memo_created_at", "")
                    )
                    sessions_db[session_id] = session
                
                # 更新计数器
                if sessions_db:
                    max_session_num = max([
                        int(sid.split('_')[1]) for sid in sessions_db.keys() 
                        if sid.startswith('session_') and sid.split('_')[1].isdigit()
                    ], default=0)
                    session_counter = max_session_num
                    
                print(f"已加载 {len(sessions_db)} 个会话")
        else:
            print("未找到会话数据文件，从空数据开始")
    except Exception as e:
        print(f"加载会话数据失败: {e}")
        print("将清空会话数据库以避免错误")
        sessions_db.clear()

def load_roles():
    """从文件加载角色数据"""
    global roles_db, role_counter
    try:
        if ROLES_FILE.exists():
            with open(ROLES_FILE, 'r', encoding='utf-8') as f:
                loaded_data = json.load(f)
                
                # 将 JSON 数据反序列化为 Pydantic 模型
                for role_id, role_data in loaded_data.items():
                    # 如果数据是字符串（旧格式），跳过
                    if isinstance(role_data, str):
                        print(f"跳过旧格式数据: {role_id}")
                        continue
                        
                    # 重建 RoleTrait 对象
                    custom_traits = []
                    if "custom_traits" in role_data and role_data["custom_traits"]:
                        for trait_data in role_data["custom_traits"]:
                            custom_traits.append(RoleTrait(key=trait_data["key"], value=trait_data["value"]))
                    
                    # 创建 CustomRole 对象
                    role = CustomRole(
                        id=role_data["id"],
                        name=role_data["name"],
                        input_spec=role_data["input_spec"],
                        output_spec=role_data["output_spec"],
                        custom_traits=custom_traits,
                        system_prompt_snapshot=role_data["system_prompt_snapshot"],
                        is_builtin=role_data.get("is_builtin", False),
                        created_at=role_data["created_at"],
                        updated_at=role_data["updated_at"],
                        usage_count=role_data.get("usage_count", 0),
                        last_used_at=role_data.get("last_used_at")
                    )
                    roles_db[role_id] = role
                
                # 更新计数器
                if roles_db:
                    max_role_num = max([
                        int(rid.split('_')[1]) for rid in roles_db.keys() 
                        if rid.startswith('custom_') and rid.split('_')[1].isdigit()
                    ], default=0)
                    role_counter = max_role_num
                    
                print(f"已加载 {len(roles_db)} 个自定义角色")
        else:
            print("未找到角色数据文件，从空数据开始")
    except Exception as e:
        print(f"加载角色数据失败: {e}")
        print("将清空角色数据库以避免错误")
        roles_db.clear()

def save_all_data():
    """保存所有数据"""
    save_sessions()
    save_roles()

def backup_data():
    """备份数据文件"""
    import shutil
    from datetime import datetime
    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_dir = DATA_DIR / "backups" / timestamp
        backup_dir.mkdir(parents=True, exist_ok=True)
        
        if SESSIONS_FILE.exists():
            shutil.copy2(SESSIONS_FILE, backup_dir / "sessions.json")
        if ROLES_FILE.exists():
            shutil.copy2(ROLES_FILE, backup_dir / "roles.json")
            
        print(f"数据已备份到 {backup_dir}")
        
        # 清理旧备份（保留最近10个）
        cleanup_old_backups()
        
    except Exception as e:
        print(f"数据备份失败: {e}")

def cleanup_old_backups():
    """清理旧的备份文件，保留最近10个"""
    try:
        backup_base = DATA_DIR / "backups"
        if not backup_base.exists():
            return
            
        backups = sorted([d for d in backup_base.iterdir() if d.is_dir()], reverse=True)
        
        # 删除多余的备份
        for old_backup in backups[10:]:
            shutil.rmtree(old_backup)
            print(f"已清理旧备份: {old_backup.name}")
            
    except Exception as e:
        print(f"清理备份失败: {e}")

def restore_from_backup(backup_name: str = None):
    """从备份恢复数据"""
    try:
        backup_base = DATA_DIR / "backups"
        if not backup_base.exists():
            print("未找到备份目录")
            return False
            
        if backup_name:
            backup_dir = backup_base / backup_name
        else:
            # 使用最新的备份
            backups = sorted([d for d in backup_base.iterdir() if d.is_dir()], reverse=True)
            if not backups:
                print("未找到备份文件")
                return False
            backup_dir = backups[0]
            
        if not backup_dir.exists():
            print(f"备份目录不存在: {backup_dir}")
            return False
            
        # 恢复会话数据
        backup_sessions = backup_dir / "sessions.json"
        if backup_sessions.exists():
            shutil.copy2(backup_sessions, SESSIONS_FILE)
            print("会话数据已从备份恢复")
            
        # 恢复角色数据
        backup_roles = backup_dir / "roles.json"
        if backup_roles.exists():
            shutil.copy2(backup_roles, ROLES_FILE)
            print("角色数据已从备份恢复")
            
        # 重新加载数据
        load_sessions()
        load_roles()
        
        print(f"数据已从备份 {backup_dir.name} 恢复")
        return True
        
    except Exception as e:
        print(f"从备份恢复数据失败: {e}")
        return False

class RoleTrait(BaseModel):
    """角色特征"""
    key: str
    value: str

class CustomRole(BaseModel):
    id: str
    name: str
    title: Optional[str] = ""  # 职业标题，显示在对话历史卡片中
    input_spec: str
    output_spec: str
    custom_traits: List[RoleTrait] = []  # 最多3个特征
    system_prompt_snapshot: str = ""
    is_builtin: bool = False
    created_at: str
    updated_at: str
    usage_count: int = 0
    last_used_at: Optional[str] = None

def generate_system_prompt(role_name: str, input_spec: str, output_spec: str, custom_traits: List[RoleTrait] = None) -> str:
    """根据设计方案生成增强的系统提示词"""
    
    # 角色特征部分
    traits_section = ""
    if custom_traits:
        traits_section = "\n【角色特征】\n"
        for trait in custom_traits:
            traits_section += f"• {trait.key}：{trait.value}\n"
    
    # 基于输入规范生成职责描述
    responsibilities = f"基于用户提供的内容，你需要按照专业标准处理{input_spec.split('、')[0] if '、' in input_spec else input_spec}"
    
    # 风格提取（从特征中）
    style = "专业性"
    if custom_traits:
        for trait in custom_traits:
            if "性格" in trait.key or "风格" in trait.key:
                style = trait.value
                break
    
    return f"""你是一个专业的【{role_name}】。
{traits_section}
【核心职责】
{responsibilities}

【工作流程】
1. 分析用户输入是否完整
2. 如信息不足，提出1-3个澄清问题
3. 按专业标准处理内容{f'（保持{style}风格）' if style != '专业性' else ''}
4. 依据输出要求组织回复

【输出要求】
{output_spec}

【质量标准】
- 保持{style}和一致性
- 每次回复风格统一
- 不编造不存在的信息"""

class SessionData(BaseModel):
    id: str
    title: str
    lastActiveAt: str
    lastMessagePreview: str = ""
    isFavorite: bool = False
    isArchived: bool = False
    roleId: str = "emma"
    model: str = "doubao-seed-1-6-250615"
    searchEnabled: bool = True
    messages: List[Dict] = []
    memo: str = ""  # 对话备忘录，用于长对话的上下文压缩
    memo_created_at: str = ""  # 备忘录创建时间

# 对话管理常量
CONVERSATION_THRESHOLD = 5   # 对话轮次阈值，超过后触发备忘录生成（实用性优化）
KEEP_RECENT_MESSAGES = 6     # 保留的最新消息数量（3轮对话=6条消息）

async def get_session_by_id(session_id: str) -> Optional[SessionData]:
    """根据session_id获取会话数据（内存存储版本，仅作为数据库服务的降级选项）"""
    return sessions_db.get(session_id)

async def summarize_conversation(messages: List[Dict], current_memo: str = "") -> str:
    """
    总结对话为备忘录
    参数:
        messages: 需要总结的消息列表
        current_memo: 现有备忘录内容
    返回:
        总结后的备忘录
    """
    try:
        # 构建对话文本
        conversation_text = ""
        for msg in messages:
            role = "用户" if msg["role"] == "user" else "助手"
            conversation_text += f"{role}: {msg['content']}\n"
        
        # 构建总结提示词
        summary_prompt = f"""请将以下对话总结为简洁的备忘录（500-800字），保留关键信息和上下文：

{'【现有备忘录】' + current_memo if current_memo else ''}

【需要总结的对话】
{conversation_text}

请生成一个简洁的备忘录，重点包括：
1. 对话的主要话题和问题
2. 重要的信息和结论
3. 用户的具体需求和偏好
4. 任何需要延续的上下文

要求：
- 备忘录长度控制在500-800字之间
- 语言简洁明了，便于后续对话参考
- 重点突出用户的个性化需求和偏好
- 使用结构化的段落组织信息"""
        
        # 使用模型管理器进行总结
        response = await model_manager.chat_completion([
            {"role": "user", "content": summary_prompt}
        ], model_id="doubao-seed-1-6-250615")
        
        return response.strip()
        
    except Exception as e:
        print(f"对话总结失败: {e}")
        # 返回简单的备份总结
        return f"对话备忘录（{len(messages)}轮对话）- 创建时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"

def should_update_memo(session: SessionData) -> bool:
    """
    判断是否应该创建或更新备忘录
    规则：
    1. 首次生成：达到CONVERSATION_THRESHOLD轮对话且无备忘录
    2. 定期更新：每增加CONVERSATION_THRESHOLD轮对话就更新一次
    
    注意：此函数在添加用户消息后调用，需要考虑即将形成的完整对话轮次
    """
    total_messages = len(session.messages)
    
    # 如果刚添加了用户消息，下一条AI回复将完成一轮对话
    # 所以我们检查的是 (total_messages + 1) // 2 轮对话
    upcoming_turns = (total_messages + 1) // 2
    
    if upcoming_turns < CONVERSATION_THRESHOLD:
        return False
    
    # 如果没有备忘录，即将达到阈值就创建
    if not session.memo:
        return upcoming_turns >= CONVERSATION_THRESHOLD
    
    # 如果已有备忘录，每增加阈值轮次就更新一次
    # 例如：阈值=5，则在10轮、15轮、20轮...时更新
    return upcoming_turns % CONVERSATION_THRESHOLD == 0

def get_context_messages(session: SessionData) -> List[Dict]:
    """
    获取用于AI生成的上下文消息
    如果有备忘录，返回备忘录+最近10轮对话
    否则返回所有对话
    """
    messages = session.messages
    
    if session.memo and len(messages) > KEEP_RECENT_MESSAGES:
        # 返回最近的消息
        recent_messages = messages[-KEEP_RECENT_MESSAGES:]
        
        # 在开头插入备忘录作为系统消息
        context_messages = [
            {"role": "system", "content": f"对话备忘录：{session.memo}"}
        ] + recent_messages
        
        return context_messages
    else:
        # 返回所有消息
        return messages

@app.get("/api/sessions")
async def get_sessions(
    query: str = "",
    page: int = 1,
    pageSize: int = 20,
    sort: str = "lastActiveAt",
    order: str = "desc",
    filter: str = ""
):
    """获取会话列表 - 兼容设计文档API"""
    try:        
        # 优先使用数据库服务
        if db_service and db_service.initialize():
            # 转换过滤器参数
            filter_type = "all"
            if "favorite" in filter:
                filter_type = "favorite"
            elif "archived" in filter:
                filter_type = "archived"
            
            # 使用数据库查询
            result = await db_service.get_sessions(
                query=query.strip(),
                page=page,
                page_size=pageSize,
                filter_type=filter_type
            )
            
            # 转换为API响应格式，并获取消息预览
            session_list = []
            for session in result['sessions']:
                # 获取最后一条用户消息作为预览
                last_message_preview = ""
                try:
                    messages = await db_service.get_recent_messages(session.id, 10)  # 获取最近10条消息
                    # 找到最后一条用户消息
                    for msg in reversed(messages):
                        if msg.role == "user":
                            preview = msg.content[:50]
                            if len(msg.content) > 50:
                                preview += "..."
                            last_message_preview = preview
                            break
                except Exception as e:
                    # 获取消息预览失败时使用空字符串
                    pass
                
                session_list.append({
                    "id": session.id,
                    "title": session.title,
                    "lastActiveAt": session.updated_at.isoformat(),
                    "lastMessagePreview": last_message_preview,
                    "isFavorite": session.is_favorite,
                    "isArchived": session.is_archived,
                    "roleId": session.role_id,
                    "model": session.model,
                    "searchEnabled": session.search_enabled,
                    "messageCount": 0,  # 可以后续优化计算
                    "memo": session.memo or ""
                })
            
            return {
                "success": True,
                "sessions": session_list,
                "total": result['total'],
                "page": page,
                "pageSize": pageSize,
                "totalPages": (result['total'] + pageSize - 1) // pageSize
            }
        
        # 降级到内存存储（兼容性）
        print("[WARN] 数据库服务不可用，使用内存存储")
        
        # 过滤逻辑
        filtered_sessions = []
        for session in sessions_db.values():
            # 搜索过滤
            if query and query.lower() not in session.title.lower():
                continue
            
            # 状态过滤
            if "favorite" in filter and not session.isFavorite:
                continue
            if "archived" in filter and not session.isArchived:
                continue
            
            # 默认不显示归档会话（除非明确请求归档会话）
            if "archived" not in filter and session.isArchived:
                continue
            
            filtered_sessions.append(session)
        
        # 排序
        if sort == "lastActiveAt":
            filtered_sessions.sort(key=lambda x: x.lastActiveAt, reverse=(order == "desc"))
        elif sort == "title":
            filtered_sessions.sort(key=lambda x: x.title, reverse=(order == "desc"))
        
        # 分页
        total = len(filtered_sessions)
        start_idx = (page - 1) * pageSize
        end_idx = start_idx + pageSize
        paged_sessions = filtered_sessions[start_idx:end_idx]
        
        # 转换为简化格式
        session_list = []
        for session in paged_sessions:
            session_list.append({
                "id": session.id,
                "title": session.title,
                "lastActiveAt": session.lastActiveAt,
                "lastMessagePreview": session.lastMessagePreview,
                "isFavorite": session.isFavorite,
                "isArchived": session.isArchived,
                "roleId": session.roleId  # 添加角色ID
            })
        
        return {
            "sessions": session_list,
            "pagination": {
                "page": page,
                "pageSize": pageSize,
                "total": total,
                "totalPages": (total + pageSize - 1) // pageSize if total > 0 else 0
            },
            "groups": {
                "today": [s for s in session_list if s["lastActiveAt"].startswith(datetime.now().strftime("%Y-%m-%d"))],
                "earlier": [s for s in session_list if not s["lastActiveAt"].startswith(datetime.now().strftime("%Y-%m-%d"))]
            }
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取会话列表失败: {str(e)}")

@app.post("/api/sessions")
async def create_session(request: dict):
    """创建新会话"""
    try:
        global session_counter
        session_counter += 1
        
        session_id = f"session_{session_counter}"
        current_time = datetime.now().isoformat()
        
        session = SessionData(
            id=session_id,
            title=request.get("title", "新对话"),
            lastActiveAt=current_time,
            lastMessagePreview="",
            isFavorite=False,
            isArchived=False,
            roleId=request.get("roleId", "emma"),
            model=request.get("model", "doubao-seed-1-6-250615"),
            searchEnabled=request.get("searchEnabled", True),
            messages=[]
        )
        
        sessions_db[session_id] = session
        save_sessions()  # 保存数据
        
        return {
            "success": True,
            "data": {
                "id": session.id,
                "title": session.title,
                "lastActiveAt": session.lastActiveAt,
                "lastMessagePreview": session.lastMessagePreview,
                "isFavorite": session.isFavorite,
                "isArchived": session.isArchived
            }
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"创建会话失败: {str(e)}")

@app.patch("/api/sessions/{session_id}")
async def update_session(session_id: str, request: dict):
    """更新会话"""
    try:
        # 优先使用数据库服务
        if db_service and db_service.is_available():
            # 构建更新数据
            updates = {}
            if "title" in request:
                updates["title"] = request["title"]
            if "isFavorite" in request:
                updates["is_favorite"] = request["isFavorite"]
            if "isArchived" in request:
                updates["is_archived"] = request["isArchived"]
            if "roleId" in request:
                updates["role_id"] = request["roleId"]
            if "model" in request:
                updates["model"] = request["model"]
            if "searchEnabled" in request:
                updates["search_enabled"] = request["searchEnabled"]
                
            # 更新数据库
            updated_session = await db_service.update_session(session_id, updates)
            if not updated_session:
                raise HTTPException(status_code=404, detail="会话不存在")
                
            return {
                "success": True,
                "data": {
                    "id": updated_session.id,
                    "title": updated_session.title,
                    "lastActiveAt": updated_session.updated_at.isoformat(),
                    "lastMessagePreview": "",  # 可以后续优化
                    "isFavorite": updated_session.is_favorite,
                    "isArchived": updated_session.is_archived
                }
            }
        else:
            # 降级到内存存储
            if session_id not in sessions_db:
                raise HTTPException(status_code=404, detail="会话不存在")
            
            session = sessions_db[session_id]
            
            # 更新字段
            if "title" in request:
                session.title = request["title"]
            if "isFavorite" in request:
                session.isFavorite = request["isFavorite"]
            if "isArchived" in request:
                session.isArchived = request["isArchived"]
            if "roleId" in request:
                session.roleId = request["roleId"]
            if "model" in request:
                session.model = request["model"]
            if "searchEnabled" in request:
                session.searchEnabled = request["searchEnabled"]
            
            session.lastActiveAt = datetime.now().isoformat()
            save_sessions()  # 保存数据
            
            return {
                "success": True,
                "data": {
                    "id": session.id,
                    "title": session.title,
                    "lastActiveAt": session.lastActiveAt,
                    "lastMessagePreview": session.lastMessagePreview,
                    "isFavorite": session.isFavorite,
                    "isArchived": session.isArchived
                }
            }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"更新会话失败: {str(e)}")

@app.delete("/api/sessions/{session_id}")
async def delete_session(session_id: str):
    """删除会话 - 软删除"""
    try:
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        # 软删除 - 标记为归档状态
        session = sessions_db[session_id]
        session.isArchived = True
        session.lastActiveAt = datetime.now().isoformat()
        
        save_sessions()  # 保存删除状态
        
        return {
            "success": True,
            "message": "会话已删除"
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"删除会话失败: {str(e)}")

@app.post("/api/sessions/{session_id}/restore")
async def restore_session(session_id: str):
    """恢复会话"""
    try:
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        session.isArchived = False
        session.lastActiveAt = datetime.now().isoformat()
        
        return {
            "success": True,
            "data": {
                "id": session.id,
                "title": session.title,
                "lastActiveAt": session.lastActiveAt,
                "lastMessagePreview": session.lastMessagePreview,
                "isFavorite": session.isFavorite,
                "isArchived": session.isArchived
            }
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"恢复会话失败: {str(e)}")

@app.get("/api/sessions/{session_id}")
async def get_session(session_id: str):
    """获取指定会话详情"""
    try:
        # 优先使用数据库服务
        if db_service and db_service.initialize():
            session = await db_service.get_session(session_id)
            if not session:
                raise HTTPException(status_code=404, detail="会话不存在")
            
            # 获取会话消息
            messages = await db_service.get_messages(session_id)
            message_list = []
            for msg in messages:
                message_list.append({
                    "id": msg.id,
                    "role": msg.role,
                    "content": msg.content,
                    "model": msg.model,
                    "created_at": msg.created_at.isoformat(),
                    "message_index": msg.message_index
                })
            
            return {
                "success": True,
                "data": {
                    "id": session.id,
                    "title": session.title,
                    "roleId": session.role_id,
                    "model": session.model,
                    "searchEnabled": session.search_enabled,
                    "messages": message_list,
                    "lastActiveAt": session.updated_at.isoformat(),
                    "isFavorite": session.is_favorite,
                    "isArchived": session.is_archived,
                    "memo": session.memo or "",
                    "memo_created_at": session.updated_at.isoformat() if session.memo else None
                }
            }
        
        # 降级到内存存储（兼容性）
        print("[WARN] 数据库服务不可用，使用内存存储")
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        
        return {
            "success": True,
            "data": {
                "id": session.id,
                "title": session.title,
                "roleId": session.roleId,
                "model": session.model,
                "searchEnabled": session.searchEnabled,
                "messages": session.messages,
                "lastActiveAt": session.lastActiveAt,
                "isFavorite": session.isFavorite,
                "isArchived": session.isArchived,
                "memo": session.memo,
                "memo_created_at": session.memo_created_at
            }
        }
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取会话详情失败: {str(e)}")

@app.post("/api/sessions/{session_id}/messages")
async def add_message_to_session(session_id: str, request: dict):
    """向会话添加消息"""
    try:
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        message = {
            "id": f"msg_{len(session.messages) + 1}",
            "role": request["role"],
            "content": request["content"],
            "timestamp": datetime.now().isoformat()
        }
        
        session.messages.append(message)
        session.lastActiveAt = datetime.now().isoformat()
        
        # 更新消息预览
        if request["role"] == "user":
            preview = request["content"][:50]
            if len(request["content"]) > 50:
                preview += "..."
            session.lastMessagePreview = preview
            
            # 检查是否需要生成或更新备忘录（仅在添加用户消息时检查）
            if should_update_memo(session):
                try:
                    total_messages = len(session.messages)
                    upcoming_turns = (total_messages + 1) // 2
                    
                    is_first_memo = not session.memo
                    if is_first_memo:
                        print(f"会话 {session_id} 即将完成第 {upcoming_turns} 轮对话，首次生成备忘录...")
                    else:
                        print(f"会话 {session_id} 即将完成第 {upcoming_turns} 轮对话，更新现有备忘录...")
                    
                    # 获取用于总结的消息
                    # 如果总消息数少于等于KEEP_RECENT_MESSAGES，则总结前面的消息但至少保留最新2条
                    if len(session.messages) <= KEEP_RECENT_MESSAGES:
                        messages_to_summarize = session.messages[:-2] if len(session.messages) > 2 else session.messages
                    else:
                        messages_to_summarize = session.messages[:-KEEP_RECENT_MESSAGES]
                    
                    # 生成或更新备忘录（传入现有备忘录以便更新）
                    memo = await summarize_conversation(messages_to_summarize, session.memo)
                    session.memo = memo
                    session.memo_created_at = datetime.now().isoformat()
                    
                    action = "生成" if is_first_memo else "更新"
                    print(f"备忘录{action}完成，长度: {len(memo)} 字符")
                    
                except Exception as e:
                    print(f"生成备忘录失败: {e}")
        
        save_sessions()  # 保存数据
        
        return {
            "success": True,
            "data": message
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"添加消息失败: {str(e)}")

@app.get("/api/sessions/{session_id}/memo")
async def get_session_memo(session_id: str):
    """获取会话备忘录"""
    try:
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        return {
            "success": True,
            "data": {
                "memo": session.memo,
                "memo_created_at": session.memo_created_at,
                "message_count": len(session.messages),
                "user_message_count": len([msg for msg in session.messages if msg["role"] == "user"])
            }
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取备忘录失败: {str(e)}")

@app.put("/api/sessions/{session_id}/memo")
async def update_session_memo(session_id: str, request: dict):
    """更新会话备忘录"""
    try:
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        session.memo = request.get("memo", "")
        session.memo_created_at = datetime.now().isoformat()
        
        save_sessions()
        
        return {
            "success": True,
            "data": {
                "memo": session.memo,
                "memo_created_at": session.memo_created_at
            },
            "message": "备忘录更新成功"
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"更新备忘录失败: {str(e)}")

@app.post("/api/sessions/{session_id}/memo/regenerate")
async def regenerate_session_memo(session_id: str):
    """重新生成会话备忘录"""
    try:
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        
        # 检查是否有足够的对话轮次生成备忘录
        user_messages = [msg for msg in session.messages if msg["role"] == "user"]
        conversation_turns = len(user_messages)
        
        if conversation_turns < 3:  # 至少3轮对话才能生成有意义的备忘录
            return {
                "success": False,
                "message": f"对话轮次不足，当前{conversation_turns}轮，需要至少3轮对话才能生成备忘录"
            }
        
        # 获取用于总结的消息
        # 如果总消息数少于等于KEEP_RECENT_MESSAGES，则总结前面的消息但至少保留最新2条
        if len(session.messages) <= KEEP_RECENT_MESSAGES:
            messages_to_summarize = session.messages[:-2] if len(session.messages) > 2 else session.messages
        else:
            messages_to_summarize = session.messages[:-KEEP_RECENT_MESSAGES]
        
        # 重新生成备忘录
        memo = await summarize_conversation(messages_to_summarize, "")
        session.memo = memo
        session.memo_created_at = datetime.now().isoformat()
        
        save_sessions()
        
        return {
            "success": True,
            "data": {
                "memo": session.memo,
                "memo_created_at": session.memo_created_at
            },
            "message": "备忘录重新生成成功"
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"重新生成备忘录失败: {str(e)}")

# ========================= 自定义角色 API =========================

@app.get("/api/roles")
async def get_roles(include_builtin: bool = True):
    """获取角色列表（预设+自定义）"""
    try:
        # 预设角色
        builtin_roles = []
        if include_builtin:
            builtin_roles = [
                {
                    "id": "emma",
                    "name": "心理咨询师 Emma",
                    "input_spec": "情绪问题、心理疏导需求、人际关系困扰",
                    "output_spec": "温暖专业的咨询回复，提供情绪支持和心理疏导",
                    "system_prompt_snapshot": """你是Emma，一位温暖专业的心理咨询师，拥有10年临床咨询经验。

你的背景：
- 心理学博士，认知行为治疗（CBT）认证咨询师
- 专长：压力管理、情绪调节、人际关系、个人成长
- 咨询过2000+来访者，擅长倾听和共情

咨询原则：
1. 先倾听，后回应，永远不评判
2. 使用开放式问题引导自我探索
3. 识别并反映情绪
4. 提供实用的情绪管理技巧
5. 适时给予肯定和鼓励

特别提醒：
- 保持专业边界，不进行医学诊断
- 遇到危机情况，提供专业求助热线""",
                    "is_builtin": True,
                    "created_at": datetime.now().isoformat(),
                    "updated_at": datetime.now().isoformat()
                },
                {
                    "id": "sophie",
                    "name": "学术助手 Sophie",
                    "input_spec": "学术问题、论文写作、研究方法、文献综述",
                    "output_spec": "严谨的学术回复，使用学术化语言，标注引用需求",
                    "system_prompt_snapshot": """你是Sophie，一位经验丰富的学术研究助理。

你的背景：
- 博士学历，发表过20+篇论文
- 精通APA、MLA、Chicago等引用格式
- 熟悉各类学术数据库和研究方法

工作方式：
1. 帮助构建论文大纲和逻辑框架
2. 润色学术语言，让表达更专业
3. 检查引用格式和学术规范
4. 提供文献综述和研究趋势分析

输出特点：
- 使用学术化语言
- 标注[需要引用]的地方
- 提供具体的改写建议""",
                    "is_builtin": True,
                    "created_at": datetime.now().isoformat(),
                    "updated_at": datetime.now().isoformat()
                },
                {
                    "id": "mike",
                    "name": "面试教练 Mike",
                    "input_spec": "面试准备、简历优化、职业规划、薪资谈判",
                    "output_spec": "实用的面试指导，提供具体框架和可执行建议",
                    "system_prompt_snapshot": """你是Mike，一位资深HR总监和职业发展教练。

你的背景：
- 在大厂做过10年招聘
- 面试过3000+候选人
- 了解各行业的面试套路

辅导风格：
1. 直接指出问题，不绕弯子
2. 提供具体的回答框架（STAR法则等）
3. 分享真实的面试官视角
4. 给出可直接使用的回答模板

特殊技能：
- 根据JD定制面试策略
- 预测可能的面试问题
- 简历优化和亮点提炼""",
                    "is_builtin": True,
                    "created_at": datetime.now().isoformat(),
                    "updated_at": datetime.now().isoformat()
                }
            ]
        
        # 自定义角色 - 优先使用数据库服务
        custom_roles = []
        if db_service and db_service.is_available():
            try:
                db_custom_roles = await db_service.get_all_custom_roles()
                for role in db_custom_roles:
                    role_dict = {
                        "id": role.id,
                        "name": role.name,
                        "title": role.title or "",
                        "input_spec": role.input_spec,
                        "output_spec": role.output_spec,
                        "custom_traits": [{"key": trait.key, "value": trait.value} for trait in role.custom_traits] if role.custom_traits else [],
                        "system_prompt_snapshot": role.system_prompt_snapshot,
                        "is_builtin": False,  # Custom roles are never builtin
                        "created_at": role.created_at.isoformat() if hasattr(role.created_at, 'isoformat') else str(role.created_at),
                        "updated_at": role.updated_at.isoformat() if hasattr(role.updated_at, 'isoformat') else str(role.updated_at),
                        "usage_count": role.usage_count,
                        "last_used_at": role.last_used_at.isoformat() if role.last_used_at and hasattr(role.last_used_at, 'isoformat') else (str(role.last_used_at) if role.last_used_at else None)
                    }
                    custom_roles.append(role_dict)
            except Exception as e:
                # 降级到内存存储时的错误日志
                pass
                # 降级到内存存储
                for role in roles_db.values():
                    role_dict = {
                        "id": role.id,
                        "name": role.name,
                        "title": role.title or "",
                        "input_spec": role.input_spec,
                        "output_spec": role.output_spec,
                        "custom_traits": [{"key": trait.key, "value": trait.value} for trait in role.custom_traits] if role.custom_traits else [],
                        "system_prompt_snapshot": role.system_prompt_snapshot,
                        "is_builtin": role.is_builtin,
                        "created_at": role.created_at,
                        "updated_at": role.updated_at,
                        "usage_count": role.usage_count,
                        "last_used_at": role.last_used_at
                    }
                    custom_roles.append(role_dict)
        else:
            # 降级到内存存储
            for role in roles_db.values():
                role_dict = {
                    "id": role.id,
                    "name": role.name,
                    "title": role.title or "",
                    "input_spec": role.input_spec,
                    "output_spec": role.output_spec,
                    "custom_traits": [{"key": trait.key, "value": trait.value} for trait in role.custom_traits] if role.custom_traits else [],
                    "system_prompt_snapshot": role.system_prompt_snapshot,
                    "is_builtin": role.is_builtin,
                    "created_at": role.created_at,
                    "updated_at": role.updated_at,
                    "usage_count": role.usage_count,
                    "last_used_at": role.last_used_at
                }
                custom_roles.append(role_dict)
        
        return {
            "success": True,
            "builtin": builtin_roles,
            "custom": custom_roles,
            "total": len(builtin_roles) + len(custom_roles)
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取角色列表失败: {str(e)}")

@app.post("/api/roles")
async def create_custom_role(request: dict):
    """创建自定义角色"""
    try:
        global role_counter
        role_counter += 1
        
        name = request.get("name", "").strip()
        title = request.get("title", "").strip()
        input_spec = request.get("input_spec", "").strip()
        output_spec = request.get("output_spec", "").strip()
        custom_traits_data = request.get("custom_traits", [])
        
        # 验证必填字段
        if not name or not title or not input_spec or not output_spec:
            raise HTTPException(status_code=400, detail="角色名称、职业标题、输入规范和输出规范不能为空")
        
        # 名称长度验证
        if len(name) < 2 or len(name) > 20:
            raise HTTPException(status_code=400, detail="角色名称长度应在2-20字符之间")
            
        # 标题长度验证
        if len(title) < 2 or len(title) > 15:
            raise HTTPException(status_code=400, detail="职业标题长度应在2-15字符之间")
        
        # 规范长度验证
        if len(input_spec) > 800 or len(output_spec) > 800:
            raise HTTPException(status_code=400, detail="输入/输出规范不能超过800字符")
        
        # 角色特征验证
        if len(custom_traits_data) > 3:
            raise HTTPException(status_code=400, detail="最多只能添加3个角色特征")
        
        # 检查名称是否重复
        for role in roles_db.values():
            if role.name == name:
                raise HTTPException(status_code=400, detail="角色名称已存在，请使用不同的名称")
        
        # 处理角色特征
        custom_traits = []
        for trait_data in custom_traits_data:
            if isinstance(trait_data, dict) and "key" in trait_data and "value" in trait_data:
                key = trait_data["key"].strip()
                value = trait_data["value"].strip()
                if key and value and len(key) <= 10 and len(value) <= 50:
                    custom_traits.append(RoleTrait(key=key, value=value))
        
        role_id = f"custom_{role_counter}"
        current_time = datetime.now().isoformat()
        
        # 生成系统提示词
        system_prompt = generate_system_prompt(name, input_spec, output_spec, custom_traits)
        
        role = CustomRole(
            id=role_id,
            name=name,
            title=title,
            input_spec=input_spec,
            output_spec=output_spec,
            custom_traits=custom_traits,
            system_prompt_snapshot=system_prompt,
            is_builtin=False,
            created_at=current_time,
            updated_at=current_time
        )
        
        roles_db[role_id] = role
        save_roles()  # 保存数据
        
        return {
            "success": True,
            "data": {
                "id": role.id,
                "name": role.name,
                "title": role.title,
                "input_spec": role.input_spec,
                "output_spec": role.output_spec,
                "custom_traits": [{"key": trait.key, "value": trait.value} for trait in role.custom_traits],
                "system_prompt_snapshot": role.system_prompt_snapshot,
                "is_builtin": role.is_builtin,
                "created_at": role.created_at,
                "updated_at": role.updated_at
            },
            "message": f"成功创建角色: {name}"
        }
        
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=f"创建角色失败: {str(e)}")


@app.patch("/api/roles/{role_id}")
async def edit_custom_role(role_id: str, request: dict):
    """编辑自定义角色"""
    try:
        if role_id not in roles_db:
            raise HTTPException(status_code=404, detail="角色不存在")
        
        role = roles_db[role_id]
        if role.is_builtin:
            raise HTTPException(status_code=403, detail="不能编辑系统角色")
        
        name = request.get("name", "").strip()
        title = request.get("title", "").strip()
        input_spec = request.get("input_spec", "").strip()
        output_spec = request.get("output_spec", "").strip()
        custom_traits_data = request.get("custom_traits", [])
        
        # 验证必填字段
        if not name or not title or not input_spec or not output_spec:
            raise HTTPException(status_code=400, detail="角色名称、职业标题、输入规范和输出规范不能为空")
        
        # 名称长度验证
        if len(name) < 2 or len(name) > 20:
            raise HTTPException(status_code=400, detail="角色名称长度应在2-20字符之间")
            
        # 标题长度验证
        if len(title) < 2 or len(title) > 15:
            raise HTTPException(status_code=400, detail="职业标题长度应在2-15字符之间")
        
        # 规范长度验证
        if len(input_spec) > 800 or len(output_spec) > 800:
            raise HTTPException(status_code=400, detail="输入/输出规范不能超过800字符")
        
        # 角色特征验证
        if len(custom_traits_data) > 3:
            raise HTTPException(status_code=400, detail="最多只能添加3个角色特征")
        
        # 检查名称是否重复（排除当前角色）
        for rid, existing_role in roles_db.items():
            if rid != role_id and existing_role.name == name:
                raise HTTPException(status_code=400, detail="角色名称已存在，请使用不同的名称")
        
        # 处理角色特征
        custom_traits = []
        for trait_data in custom_traits_data:
            if isinstance(trait_data, dict) and "key" in trait_data and "value" in trait_data:
                key = trait_data["key"].strip()
                value = trait_data["value"].strip()
                if key and value and len(key) <= 10 and len(value) <= 50:
                    custom_traits.append(RoleTrait(key=key, value=value))
        
        # 生成新的系统提示词
        system_prompt = generate_system_prompt(name, input_spec, output_spec, custom_traits)
        
        # 更新角色
        role.name = name
        role.title = title
        role.input_spec = input_spec
        role.output_spec = output_spec
        role.custom_traits = custom_traits
        role.system_prompt_snapshot = system_prompt
        role.updated_at = datetime.now().isoformat()
        
        save_roles()  # 保存数据
        
        return {
            "success": True,
            "data": {
                "id": role.id,
                "name": role.name,
                "title": role.title,
                "input_spec": role.input_spec,
                "output_spec": role.output_spec,
                "custom_traits": [{"key": trait.key, "value": trait.value} for trait in role.custom_traits],
                "system_prompt_snapshot": role.system_prompt_snapshot,
                "is_builtin": role.is_builtin,
                "created_at": role.created_at,
                "updated_at": role.updated_at
            },
            "message": f"成功更新角色: {name}"
        }
        
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        print(f"编辑角色错误: {e}")
        raise HTTPException(status_code=500, detail="编辑角色失败")

@app.delete("/api/roles/{role_id}")
async def delete_custom_role(role_id: str):
    """删除自定义角色"""
    try:
        if role_id not in roles_db:
            raise HTTPException(status_code=404, detail="角色不存在")
        
        role = roles_db[role_id]
        
        # 不允许删除内置角色
        if role.is_builtin:
            raise HTTPException(status_code=403, detail="不能删除预设角色")
        
        del roles_db[role_id]
        save_roles()  # 保存数据
        
        return {
            "success": True,
            "message": f"已删除角色: {role.name}"
        }
        
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=f"删除角色失败: {str(e)}")

# ========================= 消息编辑 API =========================

@app.put("/api/sessions/{session_id}/messages/{message_id}")
async def edit_user_message(session_id: str, message_id: str, request: dict):
    """
    编辑用户消息并智能处理备忘录
    
    参数:
    - session_id: 会话ID
    - message_id: 消息ID（而非索引，更稳定）
    - request: {
        "content": "新的消息内容",
        "auto_regenerate": bool  # 是否自动重新生成AI回复
      }
    """
    try:
        # 1. 验证会话存在
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        
        # 2. 查找消息索引
        message_index = -1
        for i, msg in enumerate(session.messages):
            if msg.get("id") == message_id:
                message_index = i
                break
        
        if message_index == -1:
            raise HTTPException(status_code=404, detail="消息不存在")
        
        # 3. 验证是用户消息
        if session.messages[message_index]["role"] != "user":
            raise HTTPException(status_code=400, detail="只能编辑用户消息")
        
        # 4. 验证新内容
        new_content = request.get("content", "").strip()
        if not new_content:
            raise HTTPException(status_code=400, detail="消息内容不能为空")
        
        # 5. 记录被删除的消息数量
        deleted_count = len(session.messages) - message_index - 1
        deleted_messages = session.messages[message_index + 1:] if deleted_count > 0 else []
        
        # 6. 判断备忘录影响
        memo_status = "unchanged"
        memo_affected = False
        
        if session.memo:
            # 计算备忘录覆盖范围
            total_messages_before = len(session.messages)
            if total_messages_before > KEEP_RECENT_MESSAGES:
                memo_covers_up_to = total_messages_before - KEEP_RECENT_MESSAGES
                
                if message_index < memo_covers_up_to:
                    # 编辑的是备忘录覆盖范围内的消息
                    memo_affected = True
                    # 编辑后的消息数
                    remaining_messages = message_index + 1
                    
                    if remaining_messages > KEEP_RECENT_MESSAGES:
                        memo_status = "need_regenerate"
                    else:
                        memo_status = "need_clear"  # 消息太少
                elif deleted_count > 0:
                    # 编辑的是最近的消息，但删除了后续消息
                    remaining_messages = message_index + 1
                    if remaining_messages <= KEEP_RECENT_MESSAGES:
                        memo_status = "need_clear"
        
        # 7. 执行编辑
        session.messages[message_index]["content"] = new_content
        session.messages[message_index]["edited"] = True
        session.messages[message_index]["edited_at"] = datetime.now().isoformat()
        
        # 8. 删除后续消息
        session.messages = session.messages[:message_index + 1]
        
        # 9. 处理备忘录
        memo_regenerated = False
        if memo_status == "need_clear":
            session.memo = ""
            session.memo_created_at = ""
            print(f"清空备忘录：剩余消息数不足")
        elif memo_status == "need_regenerate":
            try:
                # 重新生成备忘录
                if len(session.messages) > KEEP_RECENT_MESSAGES:
                    messages_to_summarize = session.messages[:-KEEP_RECENT_MESSAGES]
                    print(f"重新生成备忘录，总结前{len(messages_to_summarize)}条消息")
                    session.memo = await summarize_conversation(messages_to_summarize, "")
                    session.memo_created_at = datetime.now().isoformat()
                    memo_regenerated = True
                    memo_status = "regenerated"
            except Exception as e:
                print(f"重新生成备忘录失败: {e}")
                memo_status = "regenerate_failed"
        
        # 10. 更新会话元数据
        session.lastActiveAt = datetime.now().isoformat()
        if session.messages:
            last_user_msg = next((msg for msg in reversed(session.messages) 
                                 if msg["role"] == "user"), None)
            if last_user_msg:
                preview = last_user_msg["content"][:50]
                if len(last_user_msg["content"]) > 50:
                    preview += "..."
                session.lastMessagePreview = preview
        
        # 11. 保存数据
        save_sessions()
        
        return {
            "success": True,
            "data": {
                "messages": session.messages,
                "deleted_count": deleted_count,
                "memo_status": memo_status,
                "memo_affected": memo_affected,
                "memo_regenerated": memo_regenerated,
                "current_memo": session.memo if session.memo else None
            },
            "message": f"消息已编辑，删除了{deleted_count}条后续对话"
        }
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"编辑消息失败: {str(e)}")

@app.get("/api/sessions/{session_id}/messages/{message_id}/edit-preview")
async def get_edit_preview(session_id: str, message_id: str):
    """
    获取编辑预览信息，包括将被删除的消息和备忘录影响
    """
    try:
        if session_id not in sessions_db:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session = sessions_db[session_id]
        
        # 查找消息索引
        message_index = -1
        for i, msg in enumerate(session.messages):
            if msg.get("id") == message_id:
                message_index = i
                break
        
        if message_index == -1:
            raise HTTPException(status_code=404, detail="消息不存在")
        
        # 获取将被删除的消息
        messages_to_delete = session.messages[message_index + 1:]
        
        # 判断备忘录影响
        memo_impact = "none"
        if session.memo and len(session.messages) > KEEP_RECENT_MESSAGES:
            memo_covers_up_to = len(session.messages) - KEEP_RECENT_MESSAGES
            if message_index < memo_covers_up_to:
                memo_impact = "will_regenerate"
            elif message_index + 1 <= KEEP_RECENT_MESSAGES:
                memo_impact = "will_clear"
        
        # 最多返回3条预览消息
        preview_messages = []
        for msg in messages_to_delete[:3]:
            preview_messages.append({
                "role": msg["role"],
                "content": msg["content"][:100] + ("..." if len(msg["content"]) > 100 else "")
            })
        
        return {
            "success": True,
            "data": {
                "delete_count": len(messages_to_delete),
                "preview_messages": preview_messages,
                "has_more": len(messages_to_delete) > 3,
                "memo_impact": memo_impact,
                "current_memo_exists": bool(session.memo)
            }
        }
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取预览失败: {str(e)}")

@app.post("/api/roles/test")
async def test_custom_role(request: dict):
    """测试角色（预览系统提示词）"""
    try:
        print(f"[DEBUG] 收到测试请求: {request}")
        name = request.get("name", "").strip()
        title = request.get("title", "").strip()
        input_spec = request.get("input_spec", "").strip()
        output_spec = request.get("output_spec", "").strip()
        custom_traits_data = request.get("custom_traits", [])
        test_input = request.get("input", "你好").strip()
        
        print(f"[DEBUG] 解析后的字段 - name: '{name}', title: '{title}', input_spec: '{input_spec}', output_spec: '{output_spec}'")
        
        if not name or not input_spec or not output_spec:
            print(f"[DEBUG] 验证失败 - name empty: {not name}, input_spec empty: {not input_spec}, output_spec empty: {not output_spec}")
            raise HTTPException(status_code=400, detail="角色名称、输入规范和输出规范不能为空")
        
        # 处理角色特征
        custom_traits = []
        for trait_data in custom_traits_data:
            if isinstance(trait_data, dict) and "key" in trait_data and "value" in trait_data:
                key = trait_data["key"].strip()
                value = trait_data["value"].strip()
                if key and value and len(key) <= 10 and len(value) <= 50:
                    custom_traits.append(RoleTrait(key=key, value=value))
        
        # 生成系统提示词预览
        system_prompt = generate_system_prompt(name, input_spec, output_spec, custom_traits)
        
        # 如果提供了测试输入，可以进行实际测试
        if test_input and test_input != "你好":
            try:
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": test_input}
                ]
                
                # 使用模型管理器进行测试
                response = await model_manager.chat_completion(messages, "doubao-seed-1-6-250615")
                
                return {
                    "success": True,
                    "system_prompt": system_prompt,
                    "test_output": response,
                    "token_estimate": len(system_prompt) // 4,  # 粗略估算
                    "message": "角色测试成功"
                }
            except Exception as test_error:
                return {
                    "success": True,
                    "system_prompt": system_prompt,
                    "test_output": f"测试失败: {str(test_error)}",
                    "message": "系统提示词生成成功，但AI测试失败"
                }
        
        return {
            "success": True,
            "system_prompt": system_prompt,
            "message": "系统提示词生成成功"
        }
        
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=f"测试角色失败: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    print("\n=== MindMate AI Chat Server ===")
    print("Server: http://localhost:8004")
    print("API Docs: http://localhost:8004/docs")
    print("Chat API: POST /api/chat")
    print("Test AI: GET /test/ai")
    print("Press Ctrl+C to stop")
    print("=" * 40)
    
    # 初始化数据库服务
    print("\n正在初始化数据库服务...")
    try:
        db_service = get_db_service()
        if db_service.initialize():
            print("[OK] 数据库服务初始化成功")
        else:
            print("[WARN] 数据库服务初始化失败，将使用文件存储作为备份")
            db_service = None
    except Exception as e:
        print(f"[ERROR] 数据库服务初始化失败: {e}")
        db_service = None
    
    # 启动时加载数据（作为备份）
    if not db_service:
        print("\n正在加载持久化数据...")
        load_sessions()
        load_roles()
        print("数据加载完成")
    else:
        print("使用数据库服务，跳过JSON文件加载")
    
    # 启动时创建备份
    print("创建启动备份...")
    backup_data()
    print("备份完成\n")
    
    uvicorn.run(
        app,
        host="127.0.0.1",
        port=8004,
        log_level="info",
        access_log=True
    )